package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

const (
	baseURL            = "https://www.okx.com"
	balanceEndpoint    = "/api/v5/account/balance"
	positionsEndpoint  = "/api/v5/account/positions"
	simulatedHeaderKey = "x-simulated-trading"

	// 这里可以写“默认”的密钥（仅在环境变量没设置时才会用到）
	// !!! 强烈建议：只在你本地填真实值，千万不要提交到 Git 或发给别人 !!!
	defaultAPIKey    = "acde133f-0cea-449b-9835-3c6a951ee955"
	defaultSecretKey = "24B147256DE61E3F1D3BA4A28F376A73"
	defaultPassphase = "Kali@123"
)

type okxBalanceResp struct {
	Code string `json:"code"`
	Msg  string `json:"msg"`
	Data []struct {
		TotalEq string `json:"totalEq"`
		Details []struct {
			Ccy      string `json:"ccy"`
			Eq       string `json:"eq"`
			CashBal  string `json:"cashBal"`
			AvailBal string `json:"availBal"`
		} `json:"details"`
	} `json:"data"`
}

type okxPositionsResp struct {
	Code string `json:"code"`
	Msg  string `json:"msg"`
	Data []struct {
		InstId  string `json:"instId"`
		PosSide string `json:"posSide"`
		Pos     string `json:"pos"`
		AvgPx   string `json:"avgPx"`
		Upl     string `json:"upl"`
		Lever   string `json:"lever"`
		MgnMode string `json:"mgnMode"`
	} `json:"data"`
}

func sign(secret, timestamp, method, path, body string) (string, error) {
	preHash := timestamp + strings.ToUpper(method) + path + body
	mac := hmac.New(sha256.New, []byte(secret))
	if _, err := mac.Write([]byte(preHash)); err != nil {
		return "", err
	}
	signature := mac.Sum(nil)
	return base64.StdEncoding.EncodeToString(signature), nil
}

func createHTTPClient() *http.Client {
	// 优先从环境变量读取代理配置
	proxyAddr := os.Getenv("ALL_PROXY")
	if proxyAddr == "" {
		proxyAddr = "socks5://192.168.150.1:7897" // 默认代理
	}

	if proxyAddr != "" {
		proxyURL, err := url.Parse(proxyAddr)
		if err != nil {
			fmt.Printf("解析代理URL失败: %v，将使用直连\n", err)
			return &http.Client{Timeout: 10 * time.Second}
		}

		// 创建带代理的 Transport
		transport := &http.Transport{
			Proxy: http.ProxyURL(proxyURL),
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			}).DialContext,
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		}

		return &http.Client{
			Transport: transport,
			Timeout:   10 * time.Second,
		}
	}

	// 无代理
	return &http.Client{
		Timeout: 10 * time.Second,
	}
}

func makeRequest(client *http.Client, method, path, body string) ([]byte, error) {
	apiKey := os.Getenv("OKX_API_KEY")
	secretKey := os.Getenv("OKX_SECRET_KEY")
	passphrase := os.Getenv("OKX_PASSPHRASE")

	// 如果环境变量没有配置，就退回到文件里的默认值
	if apiKey == "" {
		apiKey = defaultAPIKey
	}
	if secretKey == "" {
		secretKey = defaultSecretKey
	}
	if passphrase == "" {
		passphrase = defaultPassphase
	}

	if apiKey == "" || secretKey == "" || passphrase == "" {
		return nil, fmt.Errorf("缺少 API 配置，请在代码中填写 defaultXXX，或设置 OKX_API_KEY / OKX_SECRET_KEY / OKX_PASSPHRASE")
	}

	timestamp := time.Now().UTC().Format(time.RFC3339)

	sig, err := sign(secretKey, timestamp, method, path, body)
	if err != nil {
		return nil, err
	}

	var bodyReader io.Reader
	if body != "" {
		bodyReader = strings.NewReader(body)
	}

	req, err := http.NewRequest(method, baseURL+path, bodyReader)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("OK-ACCESS-KEY", apiKey)
	req.Header.Set("OK-ACCESS-SIGN", sig)
	req.Header.Set("OK-ACCESS-TIMESTAMP", timestamp)
	req.Header.Set("OK-ACCESS-PASSPHRASE", passphrase)
	// 模拟盘
	req.Header.Set(simulatedHeaderKey, "1")

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(data))
	}

	return data, nil
}

func fetchBalance(client *http.Client) (*okxBalanceResp, error) {
	raw, err := makeRequest(client, http.MethodGet, balanceEndpoint, "")
	if err != nil {
		return nil, err
	}
	var resp okxBalanceResp
	if err := json.Unmarshal(raw, &resp); err != nil {
		return nil, err
	}
	if resp.Code != "0" {
		return nil, fmt.Errorf("OKX 错误: code=%s msg=%s", resp.Code, resp.Msg)
	}
	return &resp, nil
}

func fetchPositions(client *http.Client) (*okxPositionsResp, error) {
	raw, err := makeRequest(client, http.MethodGet, positionsEndpoint, "")
	if err != nil {
		return nil, err
	}
	var resp okxPositionsResp
	if err := json.Unmarshal(raw, &resp); err != nil {
		return nil, err
	}
	if resp.Code != "0" {
		return nil, fmt.Errorf("OKX 错误: code=%s msg=%s", resp.Code, resp.Msg)
	}
	return &resp, nil
}

func clearScreen() {
	// ANSI：清屏并移动光标到左上角
	fmt.Print("\033[2J\033[H")
}

func printSnapshot(client *http.Client) {
	bal, err := fetchBalance(client)
	if err != nil {
		clearScreen()
		fmt.Println("获取余额失败:", err)
		return
	}
	pos, err := fetchPositions(client)
	if err != nil {
		clearScreen()
		fmt.Println("获取持仓失败:", err)
		return
	}

	clearScreen()
	now := time.Now().Format("2006-01-02 15:04:05")
	fmt.Printf("OKX 模拟盘账户概览（%s）\n", now)
	fmt.Println(strings.Repeat("=", 40))

	if len(bal.Data) > 0 {
		fmt.Println("总权益（totalEq）:", bal.Data[0].TotalEq)
		fmt.Println("\n主要币种余额:")
		for _, d := range bal.Data[0].Details {
			fmt.Printf("  %-6s  总额(eq): %-16s  可用(availBal): %-16s\n",
				d.Ccy, d.Eq, d.AvailBal)
		}
	} else {
		fmt.Println("暂无余额数据")
	}

	fmt.Println("\n持仓信息:")
	if len(pos.Data) == 0 {
		fmt.Println("  当前无持仓")
	} else {
		for _, p := range pos.Data {
			fmt.Printf("  %-12s  方向:%-4s  仓位:%-10s  开仓均价:%-12s  未实现盈亏:%-12s  杠杆:%s\n",
				p.InstId, p.PosSide, p.Pos, p.AvgPx, p.Upl, p.Lever)
		}
	}

	fmt.Println("\n提示: 程序每秒自动刷新。按 Ctrl+C 退出。")
}

func main() {
	client := createHTTPClient()

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	// 初次打印一次
	printSnapshot(client)

	for range ticker.C {
		printSnapshot(client)
	}
}